import com.sun.jna.Pointer;
import library.WintunLibrary;
import org.pcap4j.packet.*;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.TcpPort;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class TcpData {

    static ConcurrentHashMap<String, TcpPipe> pipeMap = new ConcurrentHashMap<>();

    public static void parse(IpPacket packet, Pointer pointer) {

        //通过源ip和源端口来判断同一个tcp流
        String key = getKey(packet);

        TcpPipe tcpPipe = pipeMap.get(key);
        if (tcpPipe == null) {
            //创建一个新的对象
            tcpPipe = new TcpPipe(pointer);
            //防止多线程覆盖对象
            pipeMap.putIfAbsent(key, tcpPipe);

        }
        tcpPipe.receive(packet);

    }

    private static String getKey(IpPacket packet) {
        InetAddress srcAddr = packet.getHeader().getSrcAddr();
        TcpPacket tcpPacket = (TcpPacket) packet.getPayload();
        TcpPort srcPort = tcpPacket.getHeader().getSrcPort();
        return srcAddr + ":" + srcPort.valueAsInt();
    }

    private static void remove(IpPacket packet) {
        pipeMap.remove(getKey(packet));
    }




    static class TcpPipe {

        public Pointer pointer;

        private AtomicBoolean synState = new AtomicBoolean(false);
        private AtomicInteger replySeq = new AtomicInteger(1000);
        private AtomicInteger replyAck = new AtomicInteger(0);

        private AtomicBoolean sendFin = new AtomicBoolean(false);
        private AtomicBoolean replyFin = new AtomicBoolean(false);


        //缓存数据，保证按顺序发送到channel
        static ConcurrentHashMap<Integer, TcpPacket> cacheData = new ConcurrentHashMap<>();


        private IpPacket srcIpPacket;

        private AsynchronousSocketChannel clientChannel;


        public TcpPipe(Pointer pointer) {
            this.pointer = pointer;
        }

        private void writeData(byte[] data) {

            Pointer packetPointer = WintunLibrary.INSTANCE.WintunAllocateSendPacket(this.pointer, data.length);

            if (packetPointer != null) {

                packetPointer.write(0, data, 0, data.length);

                // 发送数据包
                WintunLibrary.INSTANCE.WintunSendPacket(this.pointer, packetPointer);

            } else {
                System.err.println("Failed to allocate send packet.");
            }
        }

        public void receive(IpPacket packet) {

            TcpPacket tcpPacket = (TcpPacket) packet.getPayload();

            boolean syn = tcpPacket.getHeader().getSyn();
            boolean fin = tcpPacket.getHeader().getFin();
            boolean psh = tcpPacket.getHeader().getPsh();
            boolean ack = tcpPacket.getHeader().getAck();
            boolean rst = tcpPacket.getHeader().getRst();

            InetAddress dstAddr = packet.getHeader().getDstAddr();
            int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();

            int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();

            if(dstAddr instanceof  Inet4Address){
                InetAddress cache = IpCache.getCache(dstAddr);
                if(cache!=null){
                    //更换实际ip
                    dstAddr = cache;
                }
            }

            if(dstAddr.getAddress()[0]==(byte) 172 && dstAddr.getAddress()[1]==(byte) 29 ){
                remove(packet);
                return;
            }

            /*
            //只拦截 example.com
            if (dstAddr.getHostAddress().equals("172.29.1.25")) {
                try {
                    dstAddr = InetAddress.getByName("2606:2800:220:1:248:1893:25c8:1946");
                } catch (UnknownHostException e) {
                    throw new RuntimeException(e);
                }
            } else {
                close();
                return;
            }
             */

            if(rst){
                log("重置关闭");
                close();;
                return;
            }

            //认为结束
            if (sendFin.get() && replyFin.get()) {
                log("正常关闭");
                close();
                return;
            }

            if (syn) {
                //防止超时重发syn数据包
                if (!synState.getAndSet(true)) {

                    //一次握手
                    try {

                        //初始化一些参数
                        srcIpPacket = packet;
                        replyAck.set(sequenceNumber + 1);

                        clientChannel = AsynchronousSocketChannel.open();
                        clientChannel.connect(new InetSocketAddress(dstAddr, dstPort), tcpPacket, new CompletionHandler<Void, TcpPacket>() {

                            @Override
                            public void completed(Void result, TcpPacket attachment) {
                                //说明连接成功,回复ack(第二次握手)
                                replySyn();
                                //尝试读取数据
                                readEvent();
                            }

                            @Override
                            public void failed(Throwable exc, TcpPacket attachment) {
                                log("连接失败关闭");
                                exc.printStackTrace();
                                close();
                            }
                        });



                    } catch (Exception e) {
                        log("clientChannel关闭");
                        e.printStackTrace();
                        close();
                    }

                }

            } else {

                if(clientChannel == null ){
                    //重置连接
                    replyRst(tcpPacket);
                    close();
                    return;
                }

                if(!clientChannel.isOpen()){
                    return;
                }


                try {
                    Packet payload = tcpPacket.getPayload();

                    //有数据就发送到服务端
                    if (payload != null && payload.getRawData().length > 0) {

                        byte[] rawData = payload.getRawData();

                        if (sequenceNumber == replyAck.get()) {

                            if (sendFin.get()) {
                                // Encrypted Alert
                                replyAck.addAndGet(rawData.length);
                                replyAck();

                                return;
                            }

                            clientChannel.write(ByteBuffer.wrap(rawData), fin, new CompletionHandler<Integer, Boolean>() {
                                @Override
                                public void completed(Integer result, Boolean attachment) {

                                    replyAck.addAndGet(result + (attachment ? 1 : 0));
                                    replyAck();

                                    if (attachment) {
                                        replyFin.set(true);

                                        if (!sendFin.getAndSet(true)) {
                                            replyFin();
                                        }
                                    }


                                    TcpPacket tcpPacket_temp = cacheData.get(replyAck.get());
                                    if (tcpPacket_temp != null) {
                                        cacheData.remove(replyAck.get());

                                        if (sendFin.get()) {
                                            return;
                                        }

                                        if(tcpPacket_temp.getPayload() == null || tcpPacket_temp.getPayload().getRawData().length  ==  0){

                                            replyFin.set(true);
                                            replyAck.incrementAndGet();

                                            replyAck();

                                            if (!sendFin.getAndSet(true)) {
                                                replyFin();
                                            }
                                            return;
                                        }

                                        clientChannel.write(ByteBuffer.wrap(tcpPacket_temp.getPayload().getRawData()), tcpPacket_temp.getHeader().getFin(), this);
                                    }
                                }

                                @Override
                                public void failed(Throwable exc, Boolean attachment) {
                                    log("写入失败关闭");
                                    exc.printStackTrace();
                                    close();
                                }

                            });

                        } else if(sequenceNumber > replyAck.get()){

                            cacheData.put(sequenceNumber, tcpPacket);

                        }


                    } else {
                        if (fin) {

                            if(sequenceNumber == replyAck.get()){

                                replyFin.set(true);
                                replyAck.incrementAndGet();

                                replyAck();

                                if (!sendFin.getAndSet(true)) {
                                    replyFin();
                                }
                            }else if(sequenceNumber > replyAck.get()){

                                cacheData.put(sequenceNumber, tcpPacket);

                            }



                        } else if (ack) {
//                            if (replyAck.get() - tcpPacket.getHeader().getSequenceNumber() == 1) {
//                                //tcp Keep-Alive
//                                replyKeepAlive();
//                            }
                        }
                    }


                } catch (Exception e) {
                    log("其他异常关闭");
                    e.printStackTrace();
                    close();
                }

            }

        }



        private void readEvent() {

            ByteBuffer buffer = ByteBuffer.allocate(1024*1024);
            // 异步读取数据
            clientChannel.read(buffer, null, new CompletionHandler<Integer, Object>() {
                @Override
                public void completed(Integer bytesRead, Object attachment) {
                    if (bytesRead > 0) {
                        // 处理读取完成的数据
                        buffer.flip();
                        byte[] data = new byte[buffer.remaining()];
                        buffer.get(data);

                        try {

                            if (!replyFin.get()) {
                                replyPsh(data);

                                // 继续异步读取下一批数据
                                buffer.clear();
                                clientChannel.read(buffer, null, this);
                            }


                        } catch (Exception ex) {
                            System.out.println("ERR: " + ex.getMessage());
                            ex.printStackTrace();
                        }

                    } else {
                        log("read -1 关闭");
                        try {
                            System.out.println("client close " + clientChannel.getRemoteAddress());
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }

                        if (!sendFin.getAndSet(true)) {
                            log("read -1 关闭 发送 fin");
                            replyFin();
                        }

                        try {
                            clientChannel.close();
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }

                    }
                }

                @Override
                public void failed(Throwable exc, Object attachment) {
                    log("读取失败关闭");
                    exc.printStackTrace();
                    close();
                }
            });
        }
        private void log(String str){
            if(srcIpPacket!=null){
                InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();
                InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
                TcpPacket tcpPacket = (TcpPacket)srcIpPacket.getPayload();
                int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();
                int srcPort = tcpPacket.getHeader().getSrcPort().valueAsInt();
                System.out.printf("[srcAddr:%s,srcPort:%d]-[dstAddr:%s,dstPort:%d] %s \n",srcAddr,srcPort,dstAddr,dstPort,str);
            }

        }

        private void close() {
            log("关闭连接");
            if (clientChannel != null) {
                try {
                    clientChannel.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            if (srcIpPacket != null) {
                remove(srcIpPacket);
            }

        }

        private void replyRst(TcpPacket tcpPacket) {
            int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();
            int acknowledgmentNumber = tcpPacket.getHeader().getAcknowledgmentNumber();
            replySeq.set(acknowledgmentNumber);
            replyAck.set(sequenceNumber);
            buildIpPack(true);
        }

        public void replyKeepAlive() {

            IpPacket newIpacket = buildIpPack(false, false, false, true, false,null, replySeq.get() - 1);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }

        public void replySyn() {

            IpPacket newIpacket = buildIpPack(true, false, false, true, null);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);

            replySeq.incrementAndGet();
        }

        public void replyFin() {

            int i = replySeq.get();

            IpPacket newIpacket = buildIpPack(false, false, true, true, false,null, i);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);


            replySeq.set(i + 1);
        }

        public void replyAck() {


            IpPacket newIpacket = buildIpPack(false, false, false, true, null);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }

        public void replyPsh(byte[] data) {

            IpPacket newIpacket = buildIpPack(false, true, false, true, data);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);

            this.replySeq.getAndAdd(data.length);

        }


        public void replyAck(int seq) {

            IpPacket newIpacket = buildIpPack(false, false, false, true,false, null, seq);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }


        private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, byte[] data) {
            int seq = replySeq.get();
            return buildIpPack(syn, psh, fin, ack, false,data, seq);
        }

        private IpPacket buildIpPack(boolean rst) {
            int seq = replySeq.get();
            return buildIpPack(false, false, false, false, true,null, seq);
        }

        private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, boolean rst,byte[] data, int seq) {

            TcpPacket localtcpPacket = (TcpPacket) srcIpPacket.getPayload();
            TcpPort srcPort = localtcpPacket.getHeader().getSrcPort();
            TcpPort dstPort = localtcpPacket.getHeader().getDstPort();


            InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();

            TcpPacket.Builder tcpBuilder;
            tcpBuilder = new TcpPacket.Builder()
                    .srcPort(dstPort)
                    .dstPort(srcPort)
                    .sequenceNumber(seq)
                    .window(Short.MIN_VALUE)
                    .acknowledgmentNumber((int) this.replyAck.get())
                    .syn(syn)
                    .psh(psh)
                    .fin(fin)
                    .ack(ack)
                    .rst(rst)
                    .srcAddr(dstAddr)
                    .dstAddr(srcAddr)
                    .correctChecksumAtBuild(true)
                    .correctLengthAtBuild(true);


            if (psh) {
                UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
                unknownBuilder.rawData(data);
                tcpBuilder.payloadBuilder(unknownBuilder);
            }


            //创建一个ip包
            IpPacket newIpPacket;
            if (srcIpPacket instanceof IpV6Packet) {
                newIpPacket = new IpV6Packet.Builder()
                        .version(IpVersion.IPV6)
                        .srcAddr((Inet6Address) dstAddr)
                        .dstAddr((Inet6Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .payloadBuilder(tcpBuilder)
                        .trafficClass(IpV6SimpleTrafficClass.newInstance((byte) 0))
                        .flowLabel(IpV6SimpleFlowLabel.newInstance((byte) 0))
                        .nextHeader(IpNumber.IPV6)
                        .build();
            } else {

                newIpPacket = new IpV4Packet.Builder()
                        .version(IpVersion.IPV4)
                        .ttl((byte) 45)
                        .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                        .protocol(IpNumber.TCP)
                        .srcAddr((Inet4Address) dstAddr)
                        .dstAddr((Inet4Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .correctChecksumAtBuild(true)
                        .payloadBuilder(tcpBuilder)
                        .identification((byte) getIpId())
                        .build();
            }

            return newIpPacket;
        }
    }

    private static AtomicInteger ipId = new AtomicInteger(0);

    public static int getIpId() {
        return ipId.incrementAndGet();
    }


}
